import numpy as np
import cv2

min_w = 30
max_h = 30

# 检测线的高度
line_h = 420

# 线的偏移量
offset = 3

# 计算中心点
def center(x, y, w, h):
    x1 = int(w / 2)
    y1 = int(h / 2)
    cx = x + x1
    cy = y + y1
    return cx, cy

# 加载视频
cap = cv2.VideoCapture("C:\\Users\\86191\\Videos\\cc.mp4")
# 获取视频的宽度和高度
frame_width = int(cap.get(3))
frame_height = int(cap.get(4))
# 创建视频并保存
fourcc = cv2.VideoWriter_fourcc(*'XVID')
vm = cv2.VideoWriter('C:\\Users\\86191\\Videos\\op.avi', fourcc, 20.0, (frame_width, frame_height))

# 去背景
bgsm = cv2.createBackgroundSubtractorMOG2()

# kernel
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))

# 计数
ca = 0
# 存储车辆
cars = []

while True:
    # 读取视频帧
    ret,frame = cap.read()
    if ret == True:
        # 灰度化
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        # 去噪
        blur = cv2.GaussianBlur(gray, (3, 3), 5)
        # 去背景
        bg = bgsm.apply(blur)
        # 腐蚀 去掉里面小方块
        ero = cv2.erode(bg, kernel)
        # 膨胀
        dil = cv2.dilate(ero, kernel, iterations=1)
        # 闭运算
        closing = cv2.morphologyEx(dil, cv2.MORPH_CLOSE, kernel)
        closing = cv2.morphologyEx(closing, cv2.MORPH_CLOSE, kernel)
        # 轮廓查找
        contours, hierarchy = cv2.findContours(closing, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        cv2.line(frame, (0, line_h), (frame.shape[1], line_h), (0, 255, 0), 3)

        for i,c in enumerate(contours):
            # 获取轮廓
            x, y, w, h = cv2.boundingRect(c)

            # 筛选外界矩阵，判断是否有效车辆
            isValid = (w >= min_w and h >= max_h)
            if not isValid:
                continue
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 1)
            # 获取中心点
            c = center(x, y, w, h)
            # 添加
            cars.append(c)
            for x,y in cars:
                if y>line_h-offset and y<line_h+offset:
                    ca +=1
                    cars.remove((x,y))
                    print(ca)
        # 显示
        cv2.putText(frame, '车数:'+str(ca), (500, 60), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 1)
        cv2.imshow('video', frame)
        vm.write(frame)
    key = cv2.waitKey(40)
    if key == ord('q'):
        break

# 关闭摄像头
cap.release()
# 关闭窗口
cv2.destroyAllWindows()
